"""
数据集的ground truth格式转换
kitti及tum的格式：
https://github.com/MichaelGrupp/evo/wiki/Formats#kitti---kitti-dataset-pose-format
"""

import os
import numpy as np
import trans_func


def urban2tum(filename, out_file):
    data = np.loadtxt(filename, delimiter=',')
    print(data.shape)
    # 先清空文件
    with open(out_file, 'w') as f:
        pass

    # 找到初始位置
    init_T = data[0, 1:].reshape(3, 4)
    init_t = init_T[:, 3]
    for i in range(len(data)):
        if i % 1000 == 0:
            print(f"processing {i} / {len(data)}")
        timestamp = data[i, 0]/1e9
        T = data[i, 1:].reshape(3, 4)
        t = T[:, 3] - init_t     # 相对于初始位置的平移
        R = T[:, :3]
        q = trans_func.euler2quaternion(trans_func.rot2euler(R))

        # 保存
        with open(out_file, 'a') as f:
            f.write(f"{timestamp} {t[0]} {t[1]} {t[2]} {q[0]} {q[1]} {q[2]} {q[3]}\n")

    print("done!")

def tum2kitti(filename, out_file):
    #TODO
    pass

def kitti2tum(init_time, filename, time_file, out_file):
    data = np.loadtxt(filename, delimiter=' ')
    times = np.loadtxt(time_file, delimiter=' ')
    print(data.shape)
    print(times.shape)
    assert data.shape[0] == times.shape[0]
    # 先清空文件
    with open(out_file, 'w') as f:
        pass

    # 找到初始位置
    for i in range(len(data)):
        if i % 1000 == 0:
            print(f"processing {i} / {len(data)}")
        timestamp = times[i] + init_time
        T = data[i].reshape(3, 4)
        t = T[:, 3]     # 相对于初始位置的平移
        R = T[:, :3]
        q = trans_func.euler2quaternion(trans_func.rot2euler(R))

        # 保存
        with open(out_file, 'a') as f:
            f.write(f"{timestamp} {t[0]} {t[1]} {t[2]} {q[0]} {q[1]} {q[2]} {q[3]}\n")

    print("done!")


def openvins2tum(filename, out_file):
    data = np.loadtxt(filename, delimiter=' ')
    print(data.shape)
    # 先清空文件
    with open(out_file, 'w') as f:
        pass

    data_tum = data[1:, :8]
    # 保存
    np.savetxt(out_file, data_tum, delimiter=' ')
    print("done!")


if __name__ == '__main__':
    # # 这里需要更改数据位置
    # gt_file = "/media/daybeha/LTFM2/dataset/urban/urban32-yeouido/global_pose.csv"
    # out_file = "/media/daybeha/LTFM2/dataset/urban/urban32-yeouido/global_pose_tum.txt"
    # urban2tum(gt_file, out_file)

    # gt_file = "/media/daybeha/LTFM2/dataset/kitti/sequences/poses/08.txt"
    # time_file = "/media/daybeha/LTFM2/dataset/kitti/sequences/08/times.txt"
    # # init_time = 1317617734.939038038  # 00
    # # init_time = 1317621341.718356       # 02
    # # init_time = 1317354879.236474991  # 05
    # init_time = 1317357877.049290895  # 08
    # out_file = "/media/daybeha/LTFM2/dataset/kitti/sequences/poses/08_tum.txt"
    # kitti2tum(init_time, gt_file, time_file, out_file)

    # data = np.genfromtxt("/home/daybeha/Documents/github/dso_ws/src/results/trajectory_tum.txt", delimiter=' ')
    # np.savetxt("/home/daybeha/Documents/github/dso_ws/src/results/trajectory_tum_.txt", data, delimiter=' ')


    # path = "/home/daybeha/Documents/github/openvins_ws/src/open_vins/results/"
    path = "/home/daybeha/Documents/github/myratslam_ws/src/ratslam_ros/results/fromopenvins_without_vm/Fast/urban26/"
    source_file = path + "traj_estimate.txt"
    out_file = path + "traj_estimate_tum.txt"
    openvins2tum(source_file, out_file)


